from genutils.io import *
from genutils import handy_functions
import copy
import random
import os
import subprocess
from matplotlib.figure import Figure 
import pylab
from numpy import histogram, linspace, arange, array

class Synapse_info:
    def __init__(self,config_filename='configure',name='E'):
        self.name = name

        # pick out just the config options starting with name
        self.onset = from_config(name+'_onset',config_filename)
        self.tau = from_config(name+'_tau',config_filename)
        self.gmax = from_config(name+'_gmax',config_filename)
        self.rev = from_config(name+'_rev',config_filename)

    def neuron_print(self,file):
        ''' provided a file handle, this just prints the synapse info in a manner
            easily readable by the .hoc file used to run the NEURON simulations.'''
        file.write("%f\t%f\t%f\t%f\n" % (self.onset,self.tau,self.gmax,self.rev))

    def just_print(self):
        print self.name + "onset = " + str(self.onset)
        print self.name + "tau = " + str(self.tau)
        print self.name + "gmax = " + str(self.gmax)
        print self.name + "rev = " + str(self.rev)

    def set_from_list(self,list):
        if len(list) != 4:
            raise RuntimeError("Cannot set with list not of length 4.")

        self.onset = list[0]
        self.tau   = list[1]
        self.gmax  = list[2]
        self.rev   = list[3]

class Cell:
    ''' This class holds one or more Synapse_info instances and runs the simulations of the
        model in the NEURON simulation environment. '''
    def __init__(self,config_filename='config'):
        # read the name of the data file in from the config file
        data_filename = from_config_string('data_filename',config_filename)
        # read in the start time for the stimulation
        self.start_time = start_time = from_config('start_time',config_filename)
        # read in the end time for the window of interest (time we will simulate the model)
        self.end_time = end_time = from_config('end_time',config_filename)
        # read in the data sampling time step length
        self.data_dt = data_dt = from_config('data_dt',config_filename)
        self.t = arange(0, self.end_time-self.start_time, self.data_dt)
        self.g_pas = from_config('g_pas',config_filename)
        self.tau = from_config('tau_m',config_filename)

        data_file = open(data_filename)
        data_lines = data_file.readlines()
        # just get the window of interest
        data_lines = data_lines[int(start_time/data_dt):int(end_time/data_dt)]

        # read the data in
        self.data = []
        trials = int(from_config('trials',config_filename))
        for i in range(trials):
            # get collumns and make a Traces object instance (they are experimental repetitions)
            column_list = from_config('data_'+str(i+1),config_filename)
            self.data.append( traces(column_list,data_lines) )
        self.time = arange(0,end_time-start_time,data_dt)

        # a list of synapse objects
        self.synapses = []

        # set the name of the plot figure
        self.fig = None
  
    def set_synapses(self,parameter_filename):
        self.synapses = []
        synapse_names = from_config_string('synapse_names',parameter_filename)
        for name in synapse_names:
            self.synapses.append(Synapse_info(\
                                     config_filename=parameter_filename,\
                                     name = name))

    def set_synapses_from_list(self,list):
        if len(list) != 4*len(self.synapses):
            raise RuntimeError("Cannot set with list not of length 4*len(self.synapses).")

        for i in range(len(self.synapses)):
            self.synapses[i].onset = list[4*i+0]
            self.synapses[i].tau   = list[4*i+1]
            self.synapses[i].gmax  = list[4*i+2]
            self.synapses[i].rev   = list[4*i+3]
        
    def exercise(self):
        '''Run the NEURON simulation'''
        # the results of a simulation
        self.sim_results = {}

        # set up a list to hold synapse current vectors.
        self.svects_list = range(len(self.synapses))

        for i in range(len(self.data)):
            hp = self.data[i].holding_potential
            # for each holding potential, we'll have simulation results
            self.sim_results[hp] = {}
            # print out the file that NEURON reads in.
            print "simulating number " + str(i) + ' with holding potential of ' + str(hp) 
            self.neuron_print(i)
            # run the NEURON code.
            nrniv_spam_file = open('junk.out','w')
            subprocess.call(["nrniv", "sim_synapses.hoc"],stdout=nrniv_spam_file,stderr=nrniv_spam_file)
            nrniv_spam_file.close()
            #os.system("nrniv sim_synapses.hoc >& junk.out")

            # read in the results of the simulation.
            results = open("full_output.out")
            rlines = results.readlines()

            # read in the voltage trace from the simulation
            line_number = 0
            self.sim_results[hp]['V_trace'] = map(float,rlines[line_number].split())
            line_number += 1
            # read in the synaptic currents from the simulation
            for svect in self.svects_list:
                svect = map(float,rlines[line_number].split())
                line_number += 1
            self.sim_results[hp]['I_traces'] = copy.deepcopy(self.svects_list)
      
    def get_synapse_parameter(self,name):
        tokens = name.split('_')
        name_prefix = tokens[0]
        name_suffix = tokens[-1]
        for syn in self.synapses:
            if name_prefix in syn.name:
                return syn.__dict__[name_suffix]
        

    def update_synapse(self, name, new_value):
        synapse_parameters = ['onset','tau','gmax','rev']
        tokens = name.split('_')
        name_prefix = tokens[0]
        name_suffix = tokens[-1]
        for syn in self.synapses:
            if name_prefix in syn.name:
                for sp in synapse_parameters:
                    if name_suffix in sp:
                        syn.__dict__.update({sp:new_value})

    def neuron_print(self,data_num):
        ofile = open('neuron.out','w')
        ofile.write('%f\n'%(self.data_dt) )
        ofile.write('%f\n'%(self.end_time-self.start_time) )
        ofile.write('%f\n'%(self.g_pas) )
        ofile.write('%f\n'%(self.tau) )
        ofile.write('%f\n'%(self.data[data_num].holding_potential) )
        ofile.write('%d\n'%(len(self.synapses)) )
        for synapse in self.synapses:
            synapse.neuron_print(ofile)
        

    def plot_results(self,use_pylab=False):
        '''Use pylab to print the results of a simulation. use_pylab=True if not
                printing to a canvas in a gui application. '''
        # if this is the first time plotting, just plot.
        if self.fig == None:
            if use_pylab: # use_pylab if not sending info to a gui, uses pylab's gui
                self.fig = pylab.figure()
            else:
                self.fig = Figure()
            self.fig.set_edgecolor('#FFFFFF')
            self.fig.set_facecolor('#FFFFFF')
            self.axes_list = []
            self.lines_list = range(len(self.data))
            # make a subplot for each holding potential
            for i in range(len(self.data)):
                self.axes_list.append(self.fig.add_subplot(len(self.data),1,i+1))
                ax = self.axes_list[-1]
                hp = self.data[i].holding_potential
                # plot the data
                ax.plot(self.time,self.data[i].average)
                # plot the simulation results
                self.lines_list[i], = ax.plot(self.time,\
                                        self.sim_results[hp]['V_trace'])

                # only show the xticks on the bottom plot
                if i < len(self.data)-1:
                    ax.set_xticks([])
                ax.set_xlim(0,200) # show just the first 200 ms
                # set the yrange and yticks
                ymin = min(self.data[i].average[30:])
                ymax = max(self.data[i].average[30:])
                # ten percent padding
                ax.set_ylim(ymin-0.1*abs(ymin-ymax),ymax+0.1*abs(ymin-ymax)) 
                # have to convert floats to strings then back to floats 
                #     to get display right
                short = lambda f: "%2.1f" % (f) 
                # set 3 yticks
                ax.set_yticks(map(float,map(short,linspace(ymin,ymax,3)))) 
                
            return self.fig
        else: # the plot already exists, so just update it.
            for i in range(len(self.data)):
                hp = self.data[i].holding_potential
                self.lines_list[i].set_data(self.time,self.sim_results[hp]['V_trace'])
            self.fig.canvas.draw()

    def calculate_residuals(self, begin_index, end_index):
        #residuals = []
        difference_arrays = []
        for i in range(len(self.data)):
            hp = self.data[i].holding_potential
            data_array = array(self.data[i].average[begin_index:end_index])
            fit_array  = array(self.sim_results[hp]['V_trace'][begin_index:end_index])
            difference_arrays.append( (data_array - fit_array)**2 )
            #residuals.append( sum(difference_arrays[-1]) )
            # return the squared difference arrays, instead of the summed up residuals
        return difference_arrays
        

    def pylab_plot(self):
        '''Use pylab to print the results of a simulation.'''
        # if this is the first time plotting, just plot.
        if self.fig == None:
            self.fig = pylab.figure()
            self.fig.set_edgecolor('#FFFFFF')
            self.fig.set_facecolor('#FFFFFF')
            self.axes_list = []
            self.lines_list = range(len(self.data))
            # make a subplot for each holding potential
            for i in range(len(self.data)):
                self.axes_list.append(self.fig.add_subplot(len(self.data),1,i+1))
                ax = self.axes_list[-1]
                hp = self.data[i].holding_potential
                # plot the data
                ax.plot(self.time,self.data[i].average, 'k-', linewidth = 2)
                # plot the simulation results
                self.lines_list[i], = ax.plot(self.time,\
                                        self.sim_results[hp]['V_trace'], 'r-', linewidth = 2)

                # only show the xticks on the bottom plot
                if i < len(self.data)-1:
                    ax.set_xticks([])
                ax.set_xlim(0,200) # show just the first 200 ms
                # set the yrange and yticks
                ymin = min(self.data[i].average[30:])
                ymax = max(self.data[i].average[30:])
                ax.set_ylim(ymin-0.1*abs(ymin-ymax),ymax+0.1*abs(ymin-ymax)) # ten percent padding
                # have to convert floats to strings then back to floats to get display right
                short = lambda f: "%2.1f" % (f) 
                ax.set_yticks(map(float,map(short,linspace(ymin,ymax,3)))) # set 3 yticks


        else: # the plot already exists, so just update it.
            for i in range(len(self.data)):
                hp = self.data[i].holding_potential
                self.lines_list[i].set_data(self.time,self.sim_results[hp]['V_trace'])




    def gnuplot_print_sparse(self,factor=10):
        '''Print out simulation results and experimental data for gnuplot sparsely.  Only
            print out every factor(th) point.'''
        ofile = open("gnuplot.out", 'w')
        count = 0
        for ti in range(len(self.time)):
            count += 1
            # only print out 1/factor of the data points (equally spaced)
            if count%factor:
                ofile.write("%f\t" % (self.time[ti]))
                # print out the simulation results (both V and I)
                for hpi in range(len(self.holding_potentials)):
                    hp = self.data[i].holding_potential
                    ofile.write("%f\t%s\t" % (self.data[hpi].average[ti],\
                        self.sim_results[hp]['V_trace'][ti]))
                    for svi in range(len(self.sim_results[hp]['I_traces'])):
                        ofile.write("%s\t" % (self.sim_results[hp]['I_traces'][svi][ti]))
                ofile.write("\n")

    def gnuplot_print(self):
        '''Print out simulation results and experimental data for gnuplot'''
        ofile = open("gnuplot.out", 'w')
        for ti in range(len(self.time)):
            ofile.write("%f\t" % (self.time[ti]))
            # print out the simulation results (both V and I)
            for hpi in range(len(self.data)):
                hp = self.data[hpi].holding_potential
                ofile.write("%f\t%s\t" % (self.data[hpi].average[ti],\
                    self.sim_results[hp]['V_trace'][ti]))
                try:
                    for svi in range(len(self.sim_results[hp]['I_traces'])):
                        ofile.write("%s\t" % (self.sim_results[hp]['I_traces'][svi][ti]))
                except KeyError:
                    pass
            ofile.write("\n")

class traces:
    ''' This class simply contains a number of data traces in a variable called traces
    and an average of those traces in a variable called average.  The constructor of this
    classs takes a list of integers which are the column numbers of the data traces and a 
    list of lines in which the data exists.'''
    def __init__(self,column_list,line_list):
        column_list = map(int,column_list)
        self.trace = []
        for i in range(len(column_list)):
            self.trace.append([])
        for line in line_list:
            tokens = line.split()
            for i in range(len(column_list)):
                self.trace[i].append(float(tokens[column_list[i]]))
        self.calc_average_trace()
        self.calc_holding_potential()
        
    def calc_average_trace(self):
        self.average = []
        for i in range(len(self.trace[0])):
            tsum = 0
            for trace in self.trace:
                tsum += trace[i]
            tsum /= len(self.trace)
            self.average.append(tsum)
        self.average = handy_functions.running_average(self.average,10) # 10 pt running avg

    def calc_holding_potential(self):
        # the holding potential is likely the most represented bin in a histogram of the data
        histo = histogram(self.average,bins=50)
        max_index = histo[0].argsort()[-1]
        self.holding_potential = histo[1][max_index]
        
        
        

"""
# handfit values here
sinfo.Eonset = 14.5
sinfo.Etau   = 3
sinfo.Egmax  = 0.006
sinfo.Erev   = 10
sinfo.Ionset = 11
sinfo.Itau   = 8.5
sinfo.Igmax  = 0.0175
sinfo.Irev   = -78.8
sinfo.gpas   = 0.00024
sinfo.tau    = 90.0
"""
